from fastapi import APIRouter
import os
from model.pmi import PMILot, PMIWafer, PMIImage, PMIFailure, PMIDevice, PMIProduct
from model.aoi import Failure
from sqlalchemy import and_, or_
import sqlalchemy
import json
import math
from sqlalchemy import func
from utils.database import SessionLocal
from datetime import timedelta
from utils.post_item import PMILotItem, PMIWaferNGItem, PMILotNGItem, PMIProductItem, PMIProductListItem, PMIProductDeleteItem
from decimal import Decimal

router = APIRouter()

root_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


@router.post('/product/all')
async def get_all_pmi_product(item:PMIProductListItem):
    with SessionLocal() as session:
        page_index = item.page_index
        page_size = item.page_size
        pmi_products = session.query(PMIProduct).filter_by(del_flag='0')
        all_product = []
        for pmi_product in pmi_products:
            device_item = {}
            device_item["product_code"] = pmi_product.product_code
            device_item["product_name"] = pmi_product.product_name
            device_item["parameter_info"] = pmi_product.parameter_info
            all_product.append(device_item)

        if not page_index:
            return {"products": all_product}
        else:
            product_per_page = all_product[(
            page_index - 1) * page_size: page_index * page_size]
        
            total_page = math.ceil(len(all_product) / page_size)
            total_num = len(all_product)

            return {"products": product_per_page, "total_page": total_page, "total_num": total_num}
    

@router.post('/product/add')
async def add_pmi_product(item: PMIProductItem):
    with SessionLocal() as session:
        product_name = item.product_name
        parameter_info = {"distance": item.distance, "set_num_zhenhen": item.set_num_zhenhen, "set_diagonal_length": item.set_diagonal_length, "set_area": item.set_area, "set_per_area": item.set_per_area}

        product = session.query(PMIProduct).filter(and_(PMIProduct.product_name==product_name, PMIProduct.del_flag=='0')).first()

        if product:
            return {"code": 1, "msg": "添加的产品已存在"}

        product = PMIProduct(
            product_name=product_name, parameter_info=json.dumps(parameter_info))
        session.add(product)
        session.commit()

        return {"code": 0, "msg": "添加 product 成功"}
    

@router.post('/product/edit')
async def edit_pmi_product(item: PMIProductItem):
    with SessionLocal() as session:
        product_name = item.product_name
        product_code = item.product_code
        parameter_info = {"distance": item.distance, "set_num_zhenhen": item.set_num_zhenhen, "set_diagonal_length": item.set_diagonal_length, "set_area": item.set_area, "set_per_area": item.set_per_area}
        product = session.query(PMIProduct).filter(and_(PMIProduct.product_code==product_code, PMIProduct.del_flag=='0')).first()

        if product:

            product.product_name = product_name
            product.parameter_info = json.dumps(parameter_info)
            session.commit()

            return {"code": 0, "msg": "修改 product 成功"}
        return {"code": 0, "msg": "修改的 product 不存在"}
    

@router.post('/product/delete')
async def delete_pmi_product(item: PMIProductDeleteItem):
    with SessionLocal() as session:
        product_code = item.product_code
       
        product = session.query(PMIProduct).filter(and_(PMIProduct.product_code==product_code, PMIProduct.del_flag=='0')).first()

        if product:

            product.del_flag = '1'
            session.commit()

            return {"code": 0, "msg": "删除 product 成功"}
        
        return {"code": 1, "msg": "删除 product 不成功"}


@router.get('/all-pmi-device')
async def get_all_pmi_device():
    with SessionLocal() as session:
        pmi_devices = session.query(PMIDevice).all()
        all_devices = []
        for pmi_devices in pmi_devices:
            device_item = {}
            device_item["device_id"] = pmi_devices.device_id
            device_item["device_code"] = pmi_devices.device_code
            all_devices.append(device_item)

        return {"devices": all_devices}


# get all lot
@router.post('/all-lots')
async def get_all_lots(item: PMILotItem):
    # 从数据库中获取所有的 lot 数据

    device_code = item.device_code
    page_index = item.page_index
    page_size = item.page_size
    # 查询关键字
    lot_name = item.lot_name

    offset_data = page_size * (page_index-1)
    res = {}

    with SessionLocal() as session:

        if lot_name:
            lot_num = session.query(func.count(PMILot.id)).filter(PMILot.lot_name.like("%" + lot_name + "%")).scalar()
            lot_paginate = session.query(PMILot).filter(and_(PMILot.lot_name.like("%" + lot_name + "%"), PMILot.del_flag=='0', PMILot.status=="COMPLETED"), PMILot.device_code==device_code).order_by(PMILot.id.desc()).offset(offset_data).limit(page_size)
        else:
            lot_num = session.query(func.count(PMILot.id)).scalar()

            lot_paginate = session.query(PMILot).filter(and_(PMILot.del_flag=='0', PMILot.device_code==device_code, PMILot.status=="COMPLETED")).order_by(PMILot.id.desc()).offset(offset_data).limit(page_size)

        lots = lot_paginate

        res["data"] = []

        for lot in lots:
            lot_item = {}
            lot_item["lot_name"] = lot.lot_name
            lot_item["lot_code"] = lot.lot_code
            lot_item["start_time"] = (lot.create_time + timedelta(hours=8)).strftime('%Y-%m-%d %H:%M:%S')
            lot_item["wafers"] = []
            lot_code = lot.lot_code

            wafers = session.query(PMIWafer).filter(and_(PMIWafer.lot_code==lot_code, PMIWafer.del_flag=='0'))
            for wafer in wafers:
                wafer_item = {}
                wafer_item["wafer_id"] = wafer.wafer_id
                wafer_item["wafer_code"] = wafer.wafer_code
                wafer_code = wafer.wafer_code

                wafer_item["image_num"] = session.query(sqlalchemy.func.count(PMIImage.id)).filter_by(wafer_code=wafer_code).scalar()
                
                wafer_item["ng_num"] = wafer_item["image_num"] - session.query(sqlalchemy.func.count(PMIImage.id)).filter(and_(PMIImage.wafer_code==wafer_code, PMIImage.detection_info=="", PMIImage.segmentation_info=="")).scalar()

                lot_item["wafers"].append(wafer_item)

            res["data"].append(lot_item)

        res["lot_num"] = lot_num

    return res


# 推理一个lot下的所有图片，并将ng图片和正确图片分开返回
@router.post('/infer-lot-ng')
async def infer_lot_ng(item: PMILotNGItem):
    lot_code = item.lot_code
    page_index = item.page_index
    ng_page_index = item.ng_page_index
    page_size = item.page_size
    keyword = item.keyword

    with SessionLocal() as session:

        wafers = session.query(PMIWafer).filter_by(lot_code=lot_code)

        ng_predict_data = []
        predict_data = []
        all_failure_num = {}

        for wafer in wafers:

            wafer_code = wafer.wafer_code

            if keyword:

                images = session.query(PMIImage).filter(and_(PMIImage.file_name.like("%" + keyword + "%"), PMIImage.wafer_code == wafer_code))
            else:
                images = session.query(PMIImage).filter_by(wafer_code=wafer_code)

            for image in images:
                item = {}
                item["wafer_id"] = wafer.wafer_id
                item["file_name"] = image.file_name
                item["file_path"] = image.file_path

                item["detection_info"] = []
                item["segmentation_info"] = []

                if all_failure_num.get(image.failure_id):
                    all_failure_num[image.failure_id] += 1
                else:
                    all_failure_num[image.failure_id] = 1

                # 表示正常图片
                if int(image.failure_id) == 1:
                    predict_data.append(item) 
                    continue

                if image.detection_info:
                    detection_info = json.loads(image.detection_info)
                    for key, value in detection_info.items():
                        for detection in value:
                            bbox = detection[0] 
                            distance = None     
                            if str(key) == "3":
                                distance = int(detection[1])
                            item["detection_info"].append([key, bbox, None, distance])      

                if image.segmentation_info:
                    segmentation_info = json.loads(image.segmentation_info)    
                    for key, value in segmentation_info.items():
                        for segmentation in value:
                            points = segmentation[0]
                            all_xy = []
                            bbox = segmentation[1]
                            area_rate = None
                            if str(key)=="3":
                                area_rate = int(Decimal(segmentation[2]).quantize(Decimal("0.01"), rounding="ROUND_HALF_UP")*100)

                            # points 指的是一张图片的所有轮廓
                            for point in points:
                                # point 指一个 pad 的轮廓
                                point_xy = []
                                for xy in point:
                                    x = xy[0][0] + bbox[0]
                                    y = xy[0][1] + bbox[1]
                                    point_xy.append([x, y])
                                all_xy.append(point_xy)
                            item["segmentation_info"].append([key, bbox, all_xy, area_rate])  

                ng_predict_data.append(item) 

        predict_data_per_page = predict_data[(
            page_index - 1) * page_size: page_index * page_size]
        ng_predict_data_per_page = ng_predict_data[(
            ng_page_index - 1) * page_size: ng_page_index * page_size]
        
        total_page = math.ceil(len(predict_data) / page_size)
        total_num = len(predict_data)

        ng_total_page = math.ceil(len(ng_predict_data) / page_size)
        ng_total_num = len(ng_predict_data)

        class_label = []
        failures = session.query(PMIFailure).all()
        for pmi_failure in failures:
            failure_id = pmi_failure.failure_id
            failure = session.query(Failure).filter_by(failure_id=failure_id).first()
            label_item = {}
            label_item["failure_name"] = failure.failure_name
            label_item["failure_id"] = failure.failure_id
            label_item["failure_num"] = all_failure_num.get(failure.failure_id, 0)
            class_label.append(label_item)

        class_label.append({"failure_name": "正常", "failure_id": 1, "failure_num": all_failure_num.get(1, 0)})

        return {"predict_data": predict_data_per_page, "ng_predict_data": ng_predict_data_per_page, "class_label": class_label, "total_page": total_page, "total_num": total_num, "ng_total_page": ng_total_page, "ng_total_num": ng_total_num}


# 推理一个wafer下的所有图片，并将ng图片和正确图片分开返回
@router.post('/infer-wafer-ng')
async def infer_wafer_ng(item: PMIWaferNGItem):
    wafer_code = item.wafer_code
    page_index = item.page_index
    ng_page_index = item.ng_page_index
    page_size = item.page_size

    with SessionLocal() as session:

        wafer = session.query(PMIWafer).filter_by(wafer_code=wafer_code).first()

        ng_predict_data = []
        predict_data = []
        all_failure_num = {}
        lot_name = None

        if wafer:

            lot_code = wafer.lot_code
            lot_name = session.query(PMILot).filter_by(lot_code=lot_code).first().lot_name

            images = session.query(PMIImage).filter_by(wafer_code=wafer_code)

            for image in images:
                item = {}
                item["wafer_id"] = wafer.wafer_id
                item["file_name"] = image.file_name
                item["file_path"] = image.file_path

                item["detection_info"] = []
                item["segmentation_info"] = []

                if all_failure_num.get(image.failure_id):
                    all_failure_num[image.failure_id] += 1
                else:
                    all_failure_num[image.failure_id] = 1


                # 表示正常图片
                if int(image.failure_id) == 1:
                    predict_data.append(item) 
                    continue

                if image.detection_info:
                    detection_info = json.loads(image.detection_info)
                    for key, value in detection_info.items():
                        for detection in value:
                            bbox = detection[0]   
                            distance = None     
                            if str(key) == "3":
                                distance = int(detection[1])
                            item["detection_info"].append([key, bbox, None, distance])

                if image.segmentation_info:
                    segmentation_info = json.loads(image.segmentation_info)    
                    for key, value in segmentation_info.items():
                        for segmentation in value:
                            points = segmentation[0]
                            all_xy = []
                            bbox = segmentation[1]
                            area_rate = None
                            if str(key) == "3":
                                area_rate = int(Decimal(segmentation[2]).quantize(Decimal("0.01"), rounding="ROUND_HALF_UP")*100)

                            # points 指的是一张图片的所有轮廓
                            for point in points:
                                # point 指一个 pad 的轮廓
                                point_xy = []
                                for xy in point:
                                    x = xy[0][0] + bbox[0]
                                    y = xy[0][1] + bbox[1]
                                    point_xy.append([x, y])
                                all_xy.append(point_xy)
                            item["segmentation_info"].append([key, bbox, all_xy, area_rate])  

                ng_predict_data.append(item) 

        predict_data_per_page = predict_data[(
            page_index - 1) * page_size: page_index * page_size]
        
        if ng_page_index:
            ng_predict_data_per_page = ng_predict_data[(
            ng_page_index - 1) * page_size: ng_page_index * page_size]
        else:
            ng_predict_data_per_page = ng_predict_data
        
        total_page = math.ceil(len(predict_data) / page_size)
        total_num = len(predict_data)

        ng_total_page = math.ceil(len(ng_predict_data) / page_size)
        ng_total_num = len(ng_predict_data)

        class_label = []
        failures = session.query(PMIFailure).all()
        for pmi_failure in failures:
            failure_id = pmi_failure.failure_id
            failure = session.query(Failure).filter_by(failure_id=failure_id).first()
            label_item = {}
            label_item["failure_name"] = failure.failure_name
            label_item["failure_id"] = failure.failure_id
            label_item["failure_num"] = all_failure_num.get(failure.failure_id, 0)
            class_label.append(label_item)

        class_label.append({"failure_name": "正常", "failure_id": 1, "failure_num": all_failure_num.get(1, 0)})

        return {"predict_data": predict_data_per_page, "ng_predict_data": ng_predict_data_per_page, "class_label": class_label, "total_page": total_page, "total_num": total_num, "ng_total_page": ng_total_page, "ng_total_num": ng_total_num, "lot_name": lot_name}


# @router.post('/lot-count')
# def lot_count(item):
#     device_code = item.get("device_code")
#     start_time = item.get("start_time")
#     end_time = item.get("end_time")
    
#     lot = PMILot.query.filter_by(device_code=device_code).first()

#     lotName = lot.lot_name

#     wafers = PMIWafer.query.filter_by(lot_code=lot_code)

#     lot_tabel_data = []
#     wafer_tabel_data = []

#     lot_tabel_data.append({"key": "Lot Name", "value": lotName})
#     lot_tabel_data.append({"key": "Wafers", "value": wafers.count()})

#     lot_bad_dice = 0
#     lot_good_dice = 0
#     all_head_data = {}

#     for wafer in wafers:
#         wafer_map = json.loads(wafer.wafer_data)
#         wafer_index = wafer.wafer_id
#         wafer_dim = (len(wafer_map), len(wafer_map[0]))

#         # 每个 wafer 的相关统计数据
#         wafer_bad_dice = 0
#         wafer_good_dice = 0

#         all_bin_dice = {}

#         failures = PMIFailure.query.filter_by(del_flag='0')

#         for pmi_failure in failures:
#             failure_id = pmi_failure.failure_id
#             failure = Failure.query.filter_by(failure_id=failure_id).first()
#             key = "bin%s_dice" % failure.failure_id
#             all_bin_dice[key] = 0
#             all_head_data[key] = failure.failure_name

#         for i in range(wafer_dim[0]):
#             for j in range(wafer_dim[1]):
#                 failure_value = wafer_map[i][j]
#                 if failure_value == 1:
#                     wafer_good_dice += 1
#                 elif failure_value == 0:
#                     continue
#                 else:
#                     wafer_bad_dice += 1
#                     for failure_item in failure_value:
#                         fault_type = failure_item["fault_type"]

#                         # for failure in failures:
#                         #     if fault_type == failure.failure_id:
#                         #         key = "bin%s_dice" % failure.failure_id
#                         #         all_bin_dice[key] += 1
#                         #         break
#                         key = "bin%s_dice" % fault_type
#                         all_bin_dice[key] += 1

#         lot_bad_dice += wafer_bad_dice
#         lot_good_dice += wafer_good_dice

#         wafer_total_dice = wafer_bad_dice + wafer_good_dice

#         wafer_table_item = {"lot": lotName, "wafer_id": wafer_index, "total_dice": wafer_total_dice, "good_dice": wafer_good_dice, "bad_dice": wafer_bad_dice}

#         for key, value in all_bin_dice.items():
#             wafer_table_item[key] = value

#         wafer_tabel_data.append(wafer_table_item)

#     lot_total_dice = lot_good_dice + lot_bad_dice
#     lot_tabel_data.append({"key": "Total Dice", "value": lot_total_dice})
#     lot_tabel_data.append({"key": "Good Dice", "value": lot_good_dice})
#     lot_tabel_data.append({"key": "Bad Dice", "value": lot_bad_dice})

#     return jsonify({"lot_data": lot_tabel_data, "wafer_data": wafer_tabel_data, "all_head_data": all_head_data})


@router.get('/lot-delete')
async def lot_delete():
    with SessionLocal() as session:

        session.query(PMILot).delete()
        session.query(PMIWafer).delete()
        session.query(PMIImage).delete()
        session.commit()

        return "清空 pmi 数据库成功"